//#include "/afs/ihep.ac.cn/users/w/wangcw/headfile/include.h"
#include "Functions/Plotstyle.h"
#include "Functions/TerminalIO.h"
#include "RooAddPdf.h"
#include "RooChebychev.h"
#include "RooDataHist.h"
#include "RooFFTConvPdf.h"
#include "RooFitResult.h"
#include "RooGaussian.h"
#include "RooKeysPdf.h"
#include "RooPlot.h"
#include "RooRealVar.h"
#include "TArrow.h"
#include "TCanvas.h"
#include "TChain.h"
#include "TCut.h"
#include "TFile.h"
#include "TH1F.h"
#include "TLegend.h"
#include "TPaveText.h"
#include "TString.h"
using namespace std;
using namespace RooFit;
using namespace BOSS_Afterburner::Plotstyle;
using namespace BOSS_Afterburner::TerminalIO;

double find_x(double, double, double, RooRealVar, RooFFTConvPdf);

/// @addtogroup Afterburner_scripts
/// @{

/// Based on the [`Rolke` example](https://nbviewer.jupyter.org/url/root.cern.ch/doc/master/notebooks/Rolke.C.nbconvert.ipynb) that comes with ROOT, but adapted with `YamlLoader`.
/// @author   WANG Chengwei 王成伟
/// @date     May 17th, 2019
void D0_data_fit(Int_t NY, const char* filename)
{
  // Set general style options
  SetStyle();
  // Set options for "final" plots
  // SetPrelimStyle();
  // OR: create meeting style plots with stat and fitbox
  // SetMeetingStyle();
  // SetPrelimStyle();

  TFile*      FILE_I_2 = new TFile(filename);
  RooKeysPdf* sigPdf   = (RooKeysPdf*)(FILE_I_2->Get("sigPDF"));

  // ----------------create DataHist--------------------------
  TChain chain("JpsiToKpiPi0EtaAlgUserInfo");
  chain.Add("/mnt/d/IHEP/root/data/data09/Jpsi/Jpsi_dataJpsiToKpiPi0EtaAlg/"
            "ROOTFILES/Jpsi_dataJpsiToKpiPi0EtaAlg.root");

  TCut cut_Pi0     = "Pi0Mass<0.150 && Pi0Mass>0.115";
  TCut cut_Eta     = "EtaMass<0.57 && EtaMass>0.50";
  TCut cut_D0      = "D0Mass<1.95 && D0Mass>1.78";
  TCut cut_EtaChi2 = "EtaChi2<20";
  TCut cut_Pi0Chi2 = "Pi0Chi2<20";
  TCut cut_kftChi2 = "kftChi2<35";
  TCut cut         = cut_Pi0 && cut_D0 && cut_Eta && cut_EtaChi2 && cut_Pi0Chi2 && cut_kftChi2;

  //Double_t mean=1.019460;
  Double_t minRange = 1.78;
  Double_t maxRange = 1.95;

  TString    var_name = "D0Mass";
  RooRealVar x(var_name, var_name, minRange, maxRange);

  TH1F* h = new TH1F("h", "", 20, minRange, maxRange);
  chain.Draw(var_name + " >> h", cut);
  Long64_t    nentries = h->GetEntries();
  RooDataHist data("data", "data", x, h);

  // ----------------create DataSet--------------------------
  //TFile* datafile = new TFile("/besfs/groups/psip/psipgroup/user/wangcw/PhiEtap/full/Data/EEtoPhiEtapAlg_data_full.root","read");
  //TTree* datatree = (TTree*)datafile->Get("KpKmGammaPipPimAlgUserInfo");

  //RooRealVar x("kftPhiMass", "x", 0.98, 1.12);
  //RooRealVar y("kftEtapMass", "y", 0.88, 1.04);
  //RooDataSet fulldata("fulldata","",datatree,RooArgSet(x,y));
  //RooDataSet* data = fulldata.reduce("kftPhiMass>0.98 && kftPhiMass<1.12 && kftEtapMass>0.88 && kftEtapMass<1.04");

  //construct signal model with signal MC shape convolved a Gaussian
  //RooRealVar mean("mean","mean",0.0017, 0., 0.005);
  RooRealVar mean("mean", "mean", 0.00167987);
  //RooRealVar sigma("sigma","sigma",0.00004, 0., 0.0005);
  RooRealVar  sigma("sigma", "sigma", 0.0000591347);
  RooGaussian Gauss("Gauss", "Gaussian PDF", x, mean, sigma);

  RooFFTConvPdf sig("sig", "sigmc(x)gaus", x, *sigPdf, Gauss);

  //construct signal model with Breit-Wigner convolved a Gaussian
  //RooRealVar mean("mean","mean", 1.01946);
  //RooRealVar width("width","width", 0.004249, 0.001, 0.01);
  //RooRealVar sigma("sigma","sigma", 0.0001, 0., 0.001);
  //RooVoigtian sig("signal","Voigtian PDF",x,mean,width,sigma);

  //---------------- background model with Chebychev
  //RooRealVar a0("a0","", 0.1, -5., 5.);
  RooRealVar a0("a0", "", -0.131683);
  //RooRealVar a1("a1","",-0.5,-10,10);
  RooChebychev bkg("background", "background", x, RooArgList(a0));

  //---------------- background model with argus
  //RooRealVar m0("m0","", 1.2, 0., 4.);
  //RooRealVar m0("m0","", 0.987);
  //RooRealVar c("c","", -20);
  //RooRealVar p("p","", 2.);
  //RooArgusBG bkg("bkg","",x,m0,c,p);

  //---------------- background model with modified argus
  //RooRealVar p("p","p", 0.5);
  //RooRealVar b("b","b", 1.);
  //RooGenericPdf bkg("bkg","background","pow((kftPhiMass-0.98735),p)*exp(-b*(kftPhiMass-0.98735))",RooArgSet(x,p,b));

  //RooRealVar nsig("nsig","nsig",0.8*nentries,0,nentries);
  RooRealVar nsig("nsig", "nsig", 0.26 * NY);
  //RooRealVar nsig("nsig","nsig", 0.26*98);
  RooRealVar nbkg("nbkg", "nbkg", 0.2 * nentries, 0, nentries);

  RooAddPdf model("model", "model", RooArgList(sig, bkg), RooArgList(nsig, nbkg));

  auto   results = model.fitTo(data, Binning(20), Range(1.78, 1.95), Save(kTRUE), Timer(1),
                             Strategy(2), Extended(1));
  double likeall = results->minNll();
  //model->fitTo(data);

  //mean.Print();
  //sigma.Print();
  //width.Print();
  //f1.Print();
  //a0.Print();
  //a1.Print();
  nsig.Print();
  nbkg.Print();

  TCanvas* c1    = new TCanvas("c1", "c1", 700, 500);
  RooPlot* frame = x.frame(Title("; M_{K^{+}#pi^{-}#pi^{0}}(GeV/c^{2});"));
  //data.plotOn(frame);
  data.plotOn(frame, Name("data"), MarkerColor(kBlack));
  model.plotOn(frame, Components(sig), LineColor(6), LineStyle(kDashed));
  model.plotOn(frame, Components(bkg), LineColor(3), LineStyle(kDashed));
  //model.plotOn(frame);
  model.plotOn(frame, Name("model"));
  //model.plotOn(frame,Components(Gauss1),LineColor(28),LineStyle(kDashed));
  //model.plotOn(frame,Components(Gauss2),LineColor(46),LineStyle(kDashed));
  c1->SetBottomMargin(0.15);
  c1->SetLeftMargin(0.15);
  frame->GetXaxis()->CenterTitle(true);
  frame->GetXaxis()->SetTitleSize(0.06);
  frame->GetXaxis()->SetTitleOffset(1.0);
  frame->GetXaxis()->SetTitleFont(42);
  frame->GetXaxis()->SetLabelSize(0.05);
  frame->GetXaxis()->SetNdivisions(512);
  frame->GetYaxis()->CenterTitle(true);
  frame->GetYaxis()->SetTitleSize(0.06);
  frame->GetYaxis()->SetTitleOffset(1.2);
  frame->GetYaxis()->SetTitleFont(42);
  frame->GetYaxis()->SetLabelSize(0.05);
  frame->GetYaxis()->SetNdivisions(505);
  frame->SetMinimum(0.01);
  frame->Draw();
  double chisqOverNdf = frame->chiSquare("model", "data", 5);

  //------------------------double Gaussian model--------------------------//

  //double MEAN1,MEAN2,SIGMA1,SIGMA2,F1,MEAN,SIGMA;
  //MEAN1=mean1.getVal();
  //MEAN2=mean2.getVal();
  //SIGMA1=sigma1.getVal();
  //SIGMA2=sigma2.getVal();
  //F1=f1.getVal();
  //MEAN=F1*MEAN1+(1-F1)*MEAN2;
  //SIGMA=sqrt(F1*SIGMA1*SIGMA1 + (1-F1)*SIGMA2*SIGMA2);

  //-------------------------signal MC shape convolved gauss---------------------------//
  double CL    = 0.685;
  double MEAN  = find_x(0.5, minRange, maxRange, x, sig);
  double LOW   = find_x(0.5 * (1 - CL), minRange, maxRange, x, sig);
  double HIGH  = find_x(1 - 0.5 * (1 - CL), minRange, maxRange, x, sig);
  double LOWER = MEAN - 3 * (MEAN - LOW);
  double UPPER = MEAN + 3 * (HIGH - MEAN);

  cout << "UPPER: " << UPPER << endl;
  cout << "LOWER: " << LOWER << endl;
  double NSIG, NBKG, ISIG, IBKG, nbkgWindow, nbkgSideband, ESIG, EBKG, EbkgWindow, WbkgSideband,
    EbkgSideband;
  x.setRange("signal", LOWER, UPPER);
  auto isig = sig.createIntegral(x, NormSet(x), Range("signal"));
  auto ibkg = bkg.createIntegral(x, NormSet(x), Range("signal"));
  ISIG      = isig->getVal();
  IBKG      = ibkg->getVal();
  NSIG      = ISIG * (nsig.getVal());
  NBKG      = IBKG * (nbkg.getVal());
  ESIG      = sqrt(ISIG) * (nsig.getError());
  EBKG      = sqrt(IBKG) * (nbkg.getError());

  x.setRange("bkg_window", 1.819, 1.890);
  auto fracSigRange1 = bkg.createIntegral(x, x, "bkg_window");
  nbkgWindow         = nbkg.getVal() * fracSigRange1->getVal();
  EbkgWindow         = sqrt(fracSigRange1->getVal()) * (nbkg.getError());

  x.setRange("bkg_sideband", 1.06, 1.084);
  auto fracSigRange2 = bkg.createIntegral(x, x, "bkg_sideband");
  nbkgSideband       = nbkg.getVal() * fracSigRange2->getVal();
  EbkgSideband       = sqrt(fracSigRange2->getVal()) * (nbkg.getError());

  //cout << "================================================================" << endl;
  //cout << "nbkgWindow = " << nbkgWindow << endl;
  //cout << "nbkgSideband = " << nbkgSideband << endl;
  cout << "chisq/Ndf: " << chisqOverNdf << endl;
  double mineChisq = chisqOverNdf * 5;
  cout << "mineChisq: " << mineChisq << endl;
  cout << "mineNegLikelihood: " << likeall << endl;

  TPaveText* pt = new TPaveText(0.55, 0.70, 0.90, 0.90, "NDC");
  pt->SetBorderSize(1);
  pt->SetFillColor(0);
  pt->SetTextAlign(12);
  pt->SetTextFont(42);
  pt->SetTextSize(0.035);
  //pt->AddText(Form("N_{sig}^{total} = %4.1f#pm%4.1f",nsig.getVal(),nsig.getError()));
  //pt->AddText(Form("N_{bkg}^{total} = %4.0f#pm%4.0f",nbkg.getVal(),nbkg.getError()));
  pt->AddText(Form("N_{signal window}^{sig} = %4.1f #pm %4.1f", NSIG, ESIG));
  //pt->AddText(Form("N_{signal window}^{bkg} = %4.0f #pm 8",nbkgWindow));
  //pt->AddText(Form("N_{bkg} = %4.0f#pm%4.0f [#pm3#sigma]",NBKG,EBKG));
  //pt->AddText(Form("N_{sideband}^{bkg} = %4.0f #pm 9",nbkgSideband,EbkgSideband));
  pt->AddText(Form("Signal Region:  [%1.3f, %1.3f]", LOWER, UPPER));
  //pt->AddText(Form("Sideband Region:  [%1.3f, %1.3f]",1.060, 1.084));
  //pt->AddText(Form("lower = %1.3f",LOWER));
  //pt->AddText(Form("upper = %1.3f",UPPER));
  //pt->AddText(Form("ratio = %0.3f",nbkgWindow/nbkgSideband));
  pt->AddText(Form("#chi^{2}/ndof = %0.3f", chisqOverNdf));
  pt->Draw();

  auto ar1 = new TArrow(LOWER, 0.06 * nentries, LOWER, 0.010 * nentries, 0.03, "|>");
  ar1->SetFillColor(2);
  ar1->SetLineWidth(2);
  ar1->SetLineColor(2);
  ar1->Draw();
  auto ar2 = new TArrow(UPPER, 0.06 * nentries, UPPER, 0.010 * nentries, 0.03, "|>");
  ar2->SetFillColor(2);
  ar2->SetLineWidth(2);
  ar2->SetLineColor(2);
  ar2->Draw();

  auto ar3 = new TArrow(1.060, 0.06 * nentries, 1.060, 0.030 * nentries, 0.03, "|>");
  ar3->SetFillColor(28);
  ar3->SetLineWidth(2);
  ar3->SetLineColor(28);
  ar3->Draw();
  auto ar4 = new TArrow(1.084, 0.06 * nentries, 1.084, 0.030 * nentries, 0.03, "|>");
  ar4->SetFillColor(28);
  ar4->SetLineWidth(2);
  ar4->SetLineColor(28);
  ar4->Draw();

  auto leg2 = new TLegend(0.18, 0.78, 0.35, 0.88);
  leg2->SetTextSize(0.06);
  leg2->SetTextFont(42);
  leg2->SetFillColor(0);
  leg2->SetBorderSize(0);
  leg2->SetHeader("(b)");
  //leg2->Draw();
  //c1->Print("KPiPi0Eta_D0_data_fit.eps");
}

double find_x(double y, double x_low, double x_high, RooRealVar x, RooFFTConvPdf sig)
{
  if(y <= 0.) return x_low;
  if(y >= 1.) return x_high;

  double x_min = x_low;
  double x_max = x_high;

  while(1)
  {
    double x_mid;
    if((x_max - x_min) > 0.000001)
    {
      double x_mid = (x_min + x_max) / 2;
      x.setRange("signal", x_low, x_mid);
      RooAbsReal* integral = sig.createIntegral(x, NormSet(x), Range("signal"));
      double      Integral = integral->getVal();
      if(Integral == y) { return x_mid; }
      else
      {
        if(Integral > y) { x_max = x_mid; }
        else
        {
          x_min = x_mid;
        }
      }
    }
    return x_mid;
  }
}

int main(int argc, char* argv[])
{
  CheckMainArguments(argc, argv, 2);
  TString input(argv[1]);
  D0_data_fit(input.Atoi(), argv[2]);
  return 0;
}